import numpy as np
import scipy.special as sp
import matplotlib.pyplot as plt

# ====================== 参数设置 ======================
q = 1.0      # 荷载幅值
a = 0.15     # 荷载作用半径
r_max = 2*a  # 空间域最大半径
k_max = 1000.0  # 增大波数截断值
num_k = 1000   # 增加波数采样密度

# ====================== 解析解计算 ======================
def compute_Q_analytic(k, q, a):
    """解析解：Q(k) = (a*q/k) * J1(a*k)"""
    Q = np.zeros_like(k)
    mask = k != 0
    Q[mask] = (a * q / k[mask]) * sp.j1(a * k[mask])
    Q[~mask] = q * a**2 / 2  # 洛必达法则处理k=0
    return Q

# 生成波数域采样点（线性网格）
k = np.linspace(0, k_max, num_k)
Q_analytic = compute_Q_analytic(k, q, a)

# ====================== 数值正变换计算 ======================
# 生成空间域采样点（需足够密集）
num_r = 1000  # 增加空间域采样点数
r = np.linspace(0, r_max, num_r)
f_analytical = np.where(r <= a, q, 0.0)

# 数值汉克尔正变换（空间域→波数域）
Q_numerical = np.zeros_like(k)
for i in range(len(k)):
    kr = k[i] * r
    integrand = f_analytical * sp.j0(kr) * r  # 被积函数：f(r) * J0(kr) * r
    Q_numerical[i] = np.trapz(integrand, r)  # 梯形法则积分

# ====================== 数值逆变换验证 ======================
# 数值逆变换（波数域→空间域）
f_reconstructed = np.zeros_like(r)
for i in range(len(r)):
    kr = k * r[i]
    integrand = Q_analytic * sp.j0(kr) * k  # 被积函数：Q(k) * J0(kr) * k
    f_reconstructed[i] = np.trapz(integrand, k)

# ====================== 绘图对比 ======================
plt.figure(figsize=(9, 3), dpi=300)

# 空间域恢复对比
plt.subplot(1, 2, 1)
plt.plot(r, f_analytical, 'b-', linewidth=2, label='Original')
plt.plot(r, f_reconstructed, 'r--', linewidth=1.5, label='Reconstructed')
plt.xlim(0, 2*a)
plt.xlabel('Radial Position (m)', fontsize=10)
plt.ylabel('$p_r$ (Pa)', fontsize=10)
plt.title('Spatial Domain Comparison with Hankel Transform', fontsize=10)
plt.legend(loc='lower left')
plt.grid(True)

# 频域对比（对数坐标）
plt.subplot(1, 2, 2)
plt.plot(k, np.abs(Q_analytic), 'b-', label='Analytical')
plt.plot(k, np.abs(Q_numerical), 'r--', label='Numerical')
plt.xscale('log')
plt.yscale('log')
plt.xlim(1e-1, k_max)
plt.xlabel('Wavenumber (1/m)', fontsize=10)
plt.ylabel(r"|$\hat{p}_r$|", fontsize=10)
plt.title('Wavenumber Domain Comparison', fontsize=10)
plt.legend(loc='lower left')
plt.grid(True)

plt.tight_layout()
plt.savefig('Hankel_Transform.png', bbox_inches='tight')
plt.show()

# ====================== 能量守恒验证 ======================
# 空间域能量
energy_r = np.trapz(f_analytical**2 * r, r)  # ∫[f(r)]² r dr

# 波数域能量
energy_k = np.trapz(Q_analytic**2 * k, k)  # ∫[Q(k)]² k dk

print(f"Energy ratio (k/r): {energy_k / energy_r:.6f}")  # 理论值应为1.0